import os
import argparse
from PIL import Image
import numpy as np

# 添加moviepy导入
try:
    from moviepy.editor import VideoFileClip
    MOVIEPY_AVAILABLE = True
except ImportError:
    MOVIEPY_AVAILABLE = False
    print("警告: 未安装moviepy库，视频处理功能将不可用。请运行 'pip install moviepy' 安装。")


def remove_background(image, tolerance=30):
    """
    自动扣除图片背景
    """
    # 转换为RGBA模式以支持透明度
    image = image.convert("RGBA")
    
    # 获取图片数据
    data = np.array(image)
    
    # 获取左上角像素的颜色作为背景色
    bg_color = data[0, 0]
    
    # 创建掩码，标记与背景色相似的像素（在容差范围内）
    # 计算颜色差异
    diff = np.abs(data[:, :, :3].astype(np.int16) - bg_color[:3].astype(np.int16))
    mask = np.all(diff <= tolerance, axis=2)
    
    # 将背景像素设为透明
    data[mask] = [0, 0, 0, 0]
    
    # 创建新图片
    no_bg_image = Image.fromarray(data, 'RGBA')
    
    return no_bg_image


def extract_frames_from_video(video_path):
    """
    从视频中提取所有帧
    """
    if not MOVIEPY_AVAILABLE:
        print(f"错误: 无法处理视频文件 {video_path}，因为未安装moviepy库。")
        return []
    
    try:
        clip = VideoFileClip(video_path)
        frames = []
        
        # 提取所有帧
        for frame in clip.iter_frames():
            # 将numpy数组转换为PIL Image
            pil_image = Image.fromarray(frame)
            frames.append(pil_image)
        
        clip.close()  # 释放资源
        return frames
    except Exception as e:
        print(f"处理视频 {video_path} 时出错: {str(e)}")
        return []


def create_spritesheet(video_path, output_path, padding=300):
    """
    从视频创建精灵表
    video_path: 视频文件路径
    output_path: 输出精灵表的文件路径
    padding: 图片间的空隙大小
    """
    if not MOVIEPY_AVAILABLE:
        print("错误: 未安装moviepy库，无法处理视频文件。")
        return
    
    if not os.path.exists(video_path):
        print(f"错误: 视频文件不存在: {video_path}")
        return
    
    # 从视频提取帧
    video_frames = extract_frames_from_video(video_path)
    
    if not video_frames:
        print("错误: 未能从视频中提取任何帧")
        return
    
    # 去除帧的背景
    images = []
    for frame in video_frames:
        frame_no_bg = remove_background(frame)
        images.append(frame_no_bg)
    
    # 计算spritesheet的尺寸
    # 使用网格布局，尽可能接近正方形
    num_images = len(images)
    grid_size = int(np.ceil(np.sqrt(num_images)))
    
    # 计算最大宽度和高度
    max_width = max(img.width for img in images)
    max_height = max(img.height for img in images)
    
    # 添加padding到尺寸计算中
    cell_width = max_width + padding * 2
    cell_height = max_height + padding * 2
    
    # 创建spritesheet
    sheet_width = grid_size * cell_width
    sheet_height = grid_size * cell_height
    spritesheet = Image.new('RGBA', (sheet_width, sheet_height), (0, 0, 0, 0))
    
    # 将图片放置到spritesheet上
    for i, img in enumerate(images):
        row = i // grid_size
        col = i % grid_size
        
        # 计算位置（在单元格内居中放置，包含padding）
        x = col * cell_width + padding + (max_width - img.width) // 2
        y = row * cell_height + padding + (max_height - img.height) // 2
        
        # 粘贴图片
        spritesheet.paste(img, (x, y), img)
    
    # 保存spritesheet
    spritesheet.save(output_path, 'PNG')
    print(f"Spritesheet已保存到: {output_path}")
    print(f"图片网格: {grid_size}x{grid_size}, 单元格大小: {cell_width}x{cell_height}")
    print(f"从视频中提取了 {len(video_frames)} 帧")

    return output_path


if __name__ == "__main__":
    create_spritesheet(r"videos\aiji.mp4", "output.png")

